import React, { useState, useEffect } from "react";
import { Tree } from "antd";
import type { TreeProps, TreeDataNode } from "antd";
import "./style.less";
import { Icons } from "@/utils/icons";

interface TreeViewProps {
  /** Tree数据 */
  initTreeData: TreeDataNode[];
  /** 禁止选择的选项 */
  disabledKeys?: React.Key[];

  /** 已选中的选项 */
  checkedKeys?: React.Key[];

  loadData?: (treeNode: TreeDataNode) => Promise<TreeDataNode[]>;
  onCheck?: (checkedKeys: React.Key[], info: any) => void | Promise<void>; // 添加外部 onCheck 回调
}
//   {
//     title: "parent 1",
//     key: "0-0",
//     children: [
//       {
//         title: "parent 1-0",
//         key: "0-0-0",
//         disabled: true,
//         children: [
//           {
//             title: "leaf",
//             key: "0-0-0-0",
//             disableCheckbox: true,
//           },
//           {
//             title: "leaf",
//             key: "0-0-0-1",
//           },
//         ],
//       },
//       {
//         title: "parent 1-1",
//         key: "0-0-1",
//         children: [
//           {
//             title: <span style={{ color: "#1677ff" }}>sss</span>,
//             key: "0-0-1-0",
//           },
//           {
//             title: "leaf",
//             key: "0-0-1-1",
//             isLeaf: true,
//           },
//         ],
//       },
//     ],
//   },
// ];
// It's just a simple demo. You can use tree map to optimize update perf.
const updateTreeData = (
  list: TreeDataNode[],
  key: React.Key,
  children: TreeDataNode[],
  disabledKeys?: React.Key[]
): TreeDataNode[] =>
  list.map((node) => {
    const disableCheckbox = disabledKeys?.includes(node.key);
    const updatedNode = {
      ...node,
      disableCheckbox, // 确保 disableCheckbox 属性被正确更新
    };
    if (node.key === key) {
      return {
        ...updatedNode,
        children,
      };
    }
    if (node.children) {
      return {
        ...updatedNode,
        children: updateTreeData(node.children, key, children, disabledKeys),
      };
    }
    return updatedNode;
  });

const TreeView: React.FC<TreeViewProps> = ({
  initTreeData,
  loadData,
  onCheck,
  checkedKeys = [],
  disabledKeys = [],
}) => {
  const [expandedKeys, setExpandedKeys] = useState<React.Key[]>([]);
  const [selectedKeys, setSelectedKeys] = useState<React.Key[]>([]);
  const [autoExpandParent, setAutoExpandParent] = useState<boolean>(true);
  const [treeData, setTreeData] = useState(initTreeData);
  const [checkedKeysState, setCheckedKeys] = useState<
    | React.Key[]
    | {
        checked: React.Key[];
        halfChecked: React.Key[];
      }
  >([]);

  useEffect(() => {
    // 从外部传入的 checkedKeys 更新状态
    setCheckedKeys(checkedKeys);
  }, [checkedKeys]);

  useEffect(() => {
    // 从外部传入的 disabledKeys 更新状态
    setTreeData((origin: TreeDataNode[]) =>
      updateTreeData(origin, "", [], disabledKeys)
    );
  }, [disabledKeys]); // 确保 disabledKeys 的变化立即更新 treeData

  const onExpand: TreeProps["onExpand"] = (expandedKeysValue) => {
    setExpandedKeys(expandedKeysValue);
    setAutoExpandParent(false);
  };

  const handleCheck: TreeProps["onCheck"] = (checkedKeysValue, info) => {
    const result = onCheck?.(checkedKeysValue as React.Key[], info);
    if (result && typeof result.then === "function") {
      result.then((canCheckeKeys: any) => {
        setCheckedKeys(canCheckeKeys);
      });
    } else {
      setCheckedKeys(checkedKeysValue);
    }
  };

  const onSelect: TreeProps["onSelect"] = (selectedKeysValue) => {
    setSelectedKeys(selectedKeysValue);
    console.log("checkedKeysState", checkedKeysState);
  };

  const onLoadData = (node: TreeDataNode) =>
    new Promise<void>((resolve) => {
      if (node.children) {
        resolve();
        return;
      }
      if (loadData) {
        loadData(node).then((res) => {
          setTreeData((origin: TreeDataNode[]) =>
            updateTreeData(
              origin,
              node.key,
              res.map((item) => ({
                ...item,
                disableCheckbox: disabledKeys?.includes(item.key),
              })) || [],
              disabledKeys
            )
          );
          resolve();
        });
      } else {
        resolve();
      }
    });

  return (
    <Tree
      className="tree-view"
      switcherIcon={<Icons.Caretright height={16} width={16} fill="#474747" />}
      checkable
      onExpand={onExpand}
      expandedKeys={expandedKeys}
      autoExpandParent={autoExpandParent}
      onCheck={handleCheck}
      checkedKeys={checkedKeysState}
      onSelect={onSelect}
      selectedKeys={selectedKeys}
      treeData={treeData}
      loadData={onLoadData}
    />
  );
};

export default TreeView;
